from pyspark.sql import SparkSession, functions as F
from pyspark.sql.types import StructType, StringType, IntegerType

if __name__ == '__main__':
    spark = SparkSession.builder. \
        appName("test"). \
        master("local[*]"). \
        config("spark.sql.shuffle.partitions", 2). \
        getOrCreate()

    schema = StructType().add("user_id", StringType()).add("movie_id", IntegerType()).add("rank", IntegerType()).add(
        "ts", StringType())
    df = spark.read.format("csv").option("sep", "\t").option("header", False).option("encoding", "utf-8").schema(
        schema).load("/Users/cdhuangchao3/tmp/spark_demo/u.data")

    # write text，一个列，需要将df转换成单列
    df.select(F.concat_ws("---", "user_id", "movie_id", "rank", "ts")). \
        write. \
        mode("overwrite"). \
        format("text"). \
        save("../target/sql/text")

    # write csv
    df.write.mode("overwrite"). \
        format("csv"). \
        option("sep", ";"). \
        option("header", True). \
        save("../target/sql/csv")

    # write json
    df.write.mode("overwrite"). \
        format("json"). \
        save("../target/sql/json")

    # write parquet
    df.write.mode("overwrite"). \
        format("parquet"). \
        save("../target/sql/parquet")
